import cv2
import numpy as np

# 颜色配置（含亮度自适应参数）
color_config = {
    "black": {"lower": [0, 0, 0], "upper": [180, 255, 50], "v_scale": 0.5},
    "yellow": {"lower": [20, 100, 100], "upper": [30, 255, 255], "v_scale": 1.2},
    "red": {"lower": [0, 100, 100], "upper": [10, 255, 255], "v_scale": 1.0},
    "blue": {"lower": [90, 135, 20], "upper": [115, 255, 255], "v_scale": 1.0}
}


def adaptive_brightness_control(image):
    """自适应亮度预处理（自动对比度增强）"""
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    mean_brightness = np.mean(gray)

    # 根据亮度动态调整对比度
    if mean_brightness < 50:  # 低光照
        alpha = 1.5  # 对比度增强系数
        beta = 30  # 亮度提升
    elif mean_brightness > 200:  # 过曝
        alpha = 0.8
        beta = -20
    else:
        alpha = 1.0
        beta = 0
    adjusted = cv2.convertScaleAbs(image, alpha=alpha, beta=beta)
    return adjusted


def detect_colored_circles_adaptive(image):
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    result = image.copy()

    # 计算全局亮度（V通道均值）
    v_mean = np.mean(hsv[:, :, 2])

    for color_name, config in color_config.items():
        # 根据亮度调整颜色阈值（动态修改HSV范围）
        lower = np.array(config["lower"], dtype=np.uint8)
        upper = np.array(config["upper"], dtype=np.uint8)

        # 动态调整V通道阈值（示例：按比例缩放）
        v_scale = config["v_scale"]
        adjusted_lower_v = max(0, min(255, int(v_mean * v_scale)))
        adjusted_upper_v = min(255, int(v_mean * (1 + v_scale)))

        # 红色特殊处理（合并两个色相区间）
        if color_name == "red":
            lower1 = np.array([lower[0], lower[1], adjusted_lower_v], dtype=np.uint8)
            upper1 = np.array([upper[0], upper[1], adjusted_upper_v], dtype=np.uint8)
            lower2 = np.array([170, lower[1], adjusted_lower_v], dtype=np.uint8)
            upper2 = np.array([180, upper[1], adjusted_upper_v], dtype=np.uint8)
            mask1 = cv2.inRange(hsv, lower1, upper1)
            mask2 = cv2.inRange(hsv, lower2, upper2)
            mask = cv2.bitwise_or(mask1, mask2)
        else:
            lower[2] = adjusted_lower_v
            upper[2] = adjusted_upper_v
            mask = cv2.inRange(hsv, lower, upper)

        # 形态学操作（核大小根据亮度动态调整）
        kernel_size = 3 if v_mean > 100 else 5  # 低光照下增大核去噪
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

        cv2.imshow(color_name,mask)
        # 霍夫圆检测（参数动态调整）
        circles = cv2.HoughCircles(
            mask, cv2.HOUGH_GRADIENT, dp=1.2, minDist=30,
            param1=200, param2=50, minRadius=10, maxRadius=300
        )
        if circles is not None:
            circles = np.uint16(np.around(circles))
            for circle in circles[0, :]:
                x, y, r = circle[0], circle[1], circle[2]
                cv2.circle(result, (x, y), r, (0, 255, 0), 3)
                cv2.putText(result, color_name, (x - r, y - r),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
    return result


# 视频流实时处理
cap = cv2.VideoCapture(1)
while True:
    ret, frame = cap.read(1)
    if not ret:
        break

    # 自适应亮度预处理
    adjusted_frame = adaptive_brightness_control(frame)

    # 检测圆形
    result = detect_colored_circles_adaptive(adjusted_frame)

    # 显示实时亮度值
    v_mean = np.mean(cv2.cvtColor(adjusted_frame, cv2.COLOR_BGR2HSV)[:, :, 2])
    cv2.putText(result, f"Brightness: {v_mean:.1f}", (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)

    cv2.imshow("Adaptive Circle Detection", result)
    if cv2.waitKey(1) ==27:
        break

cap.release()
cv2.destroyAllWindows()